package cn.orange.ch07_binarytree;

import cn.orange.TreeNode;

import java.util.LinkedList;

/**
 * LC538.把二叉搜索树转换为累加树
 */
public class LC538 {
    public TreeNode convertBSTD(TreeNode root) { //迭代解法(累加树，从后往前遍历)
        LinkedList<TreeNode> stk = new LinkedList<>();
        TreeNode curr = root;
        TreeNode pop = null;
        int sum = 0;
        while (!stk.isEmpty() || curr != null) {
            if (curr != null) {
                stk.push(curr);
                curr = curr.right;
            } else {
                TreeNode peek = stk.peek();
                if (peek.left == null) {
                    sum += peek.val;
                    peek.val = sum;
                    pop = stk.pop();
                } else if (peek.left == pop) {
                    pop = stk.pop();
                } else {
                    sum += peek.val;
                    peek.val = sum;
                    curr = peek.left;
                }
            }
        }
        return root;
    }

    private int sum = 0;
    public TreeNode convertBST(TreeNode root) {
        traverseTree(root);
        return root;
    }
    private void traverseTree(TreeNode root) {
        if (root == null) return;
        traverseTree(root.right);
        sum += root.val;
        root.val = sum;
        traverseTree(root.left);
    }
}
